
"""
Core KL-PSS Fusion Module (Paper Section 2.2)
Implements Equations 1-4: KL divergence and PSS similarity modeling
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

class KLDistributionDivergence(nn.Module):
    """Global KL Divergence (Equation 2)"""
    def forward(self, user_dist, item_dist):
        """
        Args:
            user_dist: Global user preference Beta distribution [B,2]
            item_dist: Global item popularity Beta distribution [B,2]
        Returns:
            kl_divergence: Calculated DE_{KL}(U||I)
        """
        # Convert to Dirichlet distributions
        u_dir = F.softmax(user_dist, dim=-1) + 1e-10
        i_dir = F.softmax(item_dist, dim=-1) + 1e-10
        
        # Compute KL divergence (Equation 2)
        kl_div = (u_dir * (torch.log(u_dir) - torch.log(i_dir))).sum(dim=-1)
        return kl_div.mean()

class ProbabilitySensitiveSimilarity(nn.Module):
    """Local PSS Modeling (Equation 3)"""
    def __init__(self, gamma=0.8, tau=0.6):
        super().__init__()
        self.gamma = gamma  # Temperature coefficient
        self.tau = tau      # Adaptive threshold
        
    def forward(self, user_ratings, item_ratings):
        """
        Args:
            user_ratings: Z-normalized user ratings [B]
            item_ratings: Z-normalized item ratings [B]
        Returns:
            pss_scores: Local similarity scores [B]
        """
        # Relative difference standardization
        user_means = torch.mean(user_ratings, keepdim=True)
        user_stds = torch.std(user_ratings, keepdim=True)
        standardized = (user_ratings - user_means) / (user_stds + 1e-10)
        
        # Nonlinear threshold activation
        thresholds = torch.sigmoid(self.tau * standardized)
        
        # Final PSS calculation (Equation 3)
        pss_vals = torch.exp(-self.gamma * torch.abs(user_ratings - item_ratings))
        return pss_vals * thresholds

class DynamicFusionLayer(nn.Module):
    """Dynamic Weighted Fusion (Equation 4)"""
    def __init__(self, text_dim):
        super().__init__()
        self.text_proj = nn.Linear(text_dim, 2)
        
    def forward(self, kl_vals, pss_vals, text_features):
        """
        Args:
            kl_vals: Global KL divergence scores [B]
            pss_vals: Local PSS similarity scores [B]
            text_features: Text embedding vectors [B, D]
        Returns:
            fused_scores: Final fusion scores [B]
        """
        # Text-based weight generation (LDA weights)
        text_weights = F.softmax(self.text_proj(text_features), dim=-1)
        
        # Explicit dynamic weighting
        kl_weights = text_weights[:, 0]  # Weight for KL component
        pss_weights = text_weights[:, 1] # Weight for PSS component
        
        # Weighted fusion (Equation 4)
        fused = kl_weights * kl_vals + pss_weights * pss_vals
        return torch.sigmoid(fused)
